"""
    pth权重文件指标评估
    @author <李康>
"""

import os
from itertools import islice

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image
from sklearn.metrics import confusion_matrix
from sklearn.preprocessing import MinMaxScaler
from torch.utils.data import DataLoader
from model.deepcrack_with_uu_Sp_3 import DeepCrack_uu_Sp_3
 # 配置文件（包含模型参数、训练参数等）

# 数据预处理
from albumentations import Compose, Normalize, Resize
from albumentations.pytorch import ToTensorV2
from tqdm import tqdm

from trainer import trainer
from trainer.trainer import DeepCrackTrainer

#  多GPU模型参数转换
def adapt_model_for_multigpu(pretrained_model_):
    """
    多GPU（DataParallel）模式参数转换
    作用：将单GPU训练保存的模型参数转换为DataParallel兼容格式
    技术原理：DataParallel会自动给原始参数键名添加'module.'前缀
    典型应用场景：单GPU训练→多GPU部署时的参数格式适配
    :param pretrained_model_: 模型权重文件路径
    :return:
    """

    # pth_obj = torch.load(pretrained_model_,weights_only=False)
    pth_obj = torch.load(pretrained_model_)
    # 多GPU模式转换
    # 将单GPU模型的参数字典转换为多GPU（DataParallel）模式下的参数字典格式
    multi_gpu = True
    if multi_gpu is True:
        from collections import OrderedDict
        multi_gpu_obj = OrderedDict()
        for k, v in pth_obj.items():
            name = 'module.' + k  # 添加DataParallel前缀
            multi_gpu_obj[name] = v
        pth_obj = multi_gpu_obj
    return pth_obj


# 已训练模型加载与配置类   add by 李康
class PredModelLoader:
    """
        模型加载与配置类
        输入：
            model_pth: 模型权重文件路径  需配合model_module进行加载
            model_module: 模型模块
            use_device: 使用设备（默认使用GPU）
            device_ids: 多GPU并行计算配置（device_ids指定使用的GPU索引） 0代表第一台  若为 [0,1]就是0和1 两台
    """
    def __init__(self, model_pth, model_module, use_device=torch.device("cuda"), device_ids=None):
        # 模型加载路径
        if device_ids is None:
            device_ids = [0]
        self.pretrained_model = model_pth
        # 设置计算设备（默认优先使用GPU）
        self.device = use_device
        # 模块引入
        self.model_module = model_module
        # 并行gpu配置
        if torch.cuda.device_count() <len(device_ids):
            # 抛出异常
            raise  Exception("您的设备数量小于您配置的并行设备数，请检查您的设备数量是否正确！")
        else:
            # 多GPU并行计算配置（device_ids指定使用的GPU索引）
            self.model_module= torch.nn.DataParallel(self.model_module, device_ids=device_ids).cuda()
        # 加载权重
        self.model_module.load_state_dict(
            adapt_model_for_multigpu(self.pretrained_model),
            strict=True  # 严格匹配权重名称
        )
        self.model_module.eval()  # 设置为评估模式（关闭dropout等训练专用层）
        self.model_module.to(self.device)
        print("模型加载完成...等待测试...")
    # 获取可使用的模型     模块+pth(预训练权重)
    def get_model(self):
        """
        获取模型
        :return:
        """
        return self.model_module

# 自定义测试数据集类
class CrackTestDataset(torch.utils.data.Dataset):
    """
        自定义测试数据集类
        输入：
            img_dir: 图像文件夹路径
            mask_dir: 标签文件夹路径
            transform: 数据预处理
    """
    def __init__(self, img_dir, mask_dir, transform=None):
        self.img_dir = img_dir
        self.mask_dir = mask_dir
        self.transform = transform
        # 获取所有图像文件名
        self.images = sorted(os.listdir(img_dir))

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        """
            获取单个数据
            输入：
                idx: 数据索引
            输出：
                image: 图像数据
                mask: 标签数据
        """
        img_path = os.path.join(self.img_dir, self.images[idx])
        mask_path = os.path.join(self.mask_dir, self.images[idx].replace('.jpg', '.png'))

        image = np.array(Image.open(img_path).convert("RGB"))
        mask = np.array(Image.open(mask_path).convert("L")) / 255  # 归一化到0-1

        if self.transform:
            augmented = self.transform(image=image, mask=mask)
            image = augmented['image']
            mask = augmented['mask']

        return image, mask

    def get_image_filenames(self, start_idx, end_idx):
        """获取指定索引范围内的文件名"""
        return self.images[start_idx:end_idx]

# 指标计算函数
def calculate_metrics(y_true, y_pred):
    """
        计算指标 通过真实值和预测值计算
    :param y_true: 真实值
    :param y_pred: 预测值
    :return: 计算的指标
    """

    # #  将数据展平
    # y_true = y_true.flatten()
    # y_pred = y_pred.flatten()
    #
    # # 计算混淆矩阵
    # tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()

    # 展平并二值化预测结果（假设阈值为0.5）
    y_true = y_true.flatten().astype(int)
    y_pred = (y_pred.flatten() > 0.5).astype(int)  # 适用于概率输出

    # 显式指定labels参数确保返回2×2矩阵
    cm = confusion_matrix(y_true, y_pred, labels=[0, 1])
    tn, fp, fn, tp = cm.ravel()

    # 常用指标
    """
    基础分类指标
    准确率（Accuracy）
        作用：衡量模型整体预测正确率。
        场景：适用于类别平衡的数据集，但对不平衡数据易产生误导（如负样本占90%时，全预测负类准确率仍达90%）。
    精确率（Precision）
        作用：反映模型预测为正类的样本中实际为正类的比例，强调预测的可靠性。
        场景：FP代价高的任务（如垃圾邮件过滤，误判正常邮件为垃圾邮件代价高）。
    召回率（Recall）
        作用：衡量模型捕捉正类样本的能力，避免漏检。
        场景：FN代价高的任务（如疾病诊断，漏诊比误诊更严重）。
    特异度（Specificity）
        作用：评估模型对负类的识别能力。
        场景：需控制假警报（如金融风控中减少误判正常交易为欺诈）。
    F1分数（F1-Score）
        作用：调和精确率与召回率，适用于数据不平衡或需平衡FP/FN的场景。
        场景：综合评估模型性能（如文本分类、推荐系统）。
    """
    accuracy = (tp + tn) / (tp + tn + fp + fn) # 准确率
    precision = tp / (tp + fp + 1e-7) # 精准率
    recall = tp / (tp + fn + 1e-7) # 召回率
    specificity = tn / (tn + fp + 1e-7) # 特异度
    f1 = 2 * (precision * recall) / (precision + recall + 1e-7) # F1

    # 分割指标
    """
    交并比（IoU）
        作用：衡量预测区域与真实区域的重叠程度，直接反映分割精度。
        场景：语义分割（如自动驾驶中的物体检测），对边界敏感。
    Dice系数（Dice Coefficient）
        作用：类似IoU但更关注正类样本的匹配度，对小目标分割更敏感。
        场景：医学图像分割（如肿瘤区域小且形状不规则）。
    """
    iou = tp / (tp + fp + fn + 1e-7) # 交并比
    dice = 2 * tp / (2 * tp + fp + fn + 1e-7) # Dice

    """
        分类任务：
            优先使用精确率、召回率、F1，根据业务需求调整阈值（如医疗需高召回率，金融需高精确率）。
        分割任务：
            IoU和Dice直观反映像素级匹配度，Dice对小目标更鲁棒，IoU对边界更严格。
        不平衡数据：
            避免依赖准确率，结合F1或ROC曲线（AUC）综合评估
    """
    return {
        "Accuracy": accuracy,
        "Precision": precision,
        "Recall": recall,
        "Specificity": specificity,
        "F1": f1,
        "IoU": iou,
        "Dice": dice
    }

# 原图读取,并处理成tensor供网络输入
def input_data_processing(img_list_path,process_num=None):
    """
    输入数据预处理
    :param img_list_path: 输入图片所在文件夹
    :return:
    """
    print("输入数据预处理中...")
    if process_num is None:
        process_num = len(os.listdir(img_list_path))
    # 创建一个空的列表，用于存储处理后的图像数据
    image_tensor_list = []
    # img_list_path 遍历每一张图片
    img_list = os.listdir(img_list_path)
    img_list.sort()
    for img_name in islice(img_list,process_num):
        image_path = os.path.join(img_list_path, img_name)
        # print(image_path)
        # 图像数据预处理 变tensor+normalization
        image = cv2.imread(image_path)
        # WHC（宽度、高度、通道数） -> CWH
        image = image.transpose(2, 0, 1)
        # 归一化
        image_normalization = image / 255.0
        # 转换为numpy数组
        image_numpy = image_normalization.astype(np.float32)
        # 转换为tensor
        image_tensor = torch.from_numpy(image_numpy).float()
        # 加上batch维度
        image_tensor = image_tensor.unsqueeze(0).cuda()
        # print(image_tensor.shape)  # torch.Size([1, 3, 256, 256])
        image_tensor_list.append(image_tensor)
    return image_tensor_list

def load_label_imgs(label_list_path,process_num=None):
    """
    标签数据预处理
    :param process_num: 处理图片的数量
    :param label_list_path: 标签图片所在文件夹  这里需要image原图与label原图的数量与顺序一致
    :return:
    """
    print("标签数据预处理中...")
    # label_list_path是否为空
    if os.listdir(label_list_path)==0:
        raise ValueError("标签数据为空，请检查标签图片文件夹")
    if process_num is None:
        process_num = len(os.listdir(label_list_path))
    # 创建一个空的列表，用于存储处理后的图像数据
    label_list = os.listdir(label_list_path)
    label_list.sort()
    result_label_list = []
    for label_name in islice(label_list,process_num):
        label_path = os.path.join(label_list_path, label_name)
        label = cv2.imread(label_path)
        result_label_list.append(label)
    return result_label_list


if __name__ == "__main__":
    # ==================== 1. 模型初始化配置 ====================
    my_pretrained_model = r'../may_predict/predicted_model/DuSp3_CRKWH1001_Wight6_2_200_DuSp3_CRKWH1001_Wight6_2_200_epoch(199)_0000200_2025-05-08-13-25-51.pth'
    my_model_module = DeepCrack_uu_Sp_3()
    my_predict_model = PredModelLoader(my_pretrained_model,my_model_module).get_model()

    # ==================== 2. 输入图像数据预处理 ====================
    with torch.no_grad():# 禁用梯度计算（节省内存）  如果不禁用 导致tensor数据携带梯度 不能转numpy处理
        # 循环处理图片的数量
        MyProcessNum = 5
        # 图像数据预处理 变tensor+normalization
        input_img_list_tensor = input_data_processing('data/validation/images/', MyProcessNum)
        my_label_list = load_label_imgs('data/validation/label/', MyProcessNum)
        # 判断input_img_list_tensor与my_label_list的数量
        print('载入原图数量：', len(input_img_list_tensor))
        print('载入标签数量：', len(my_label_list))
        if len(input_img_list_tensor) != len(my_label_list):
            # print('input_img_list_tensor 与 my_label_list 数量不一致')
            raise ValueError('input_img_list_tensor 与 my_label_list 数量不一致')

        # 创建一个列表，用于存储输入图像和标签
        img_and_label_list = []
        # 存放指标的列表
        metrics = []
        for i in range(len(input_img_list_tensor)):
            # print(input_img_list_tensor[i])   处理过的原图tensor与标签array
            img_and_label_list.append({'image_tensor':input_img_list_tensor[i], 'label_tensor':my_label_list[i]})
        # print(img_and_label_list)
        for image_and_label_obj in tqdm(islice(img_and_label_list,MyProcessNum),
                                        desc='处理进度',
                                        total=MyProcessNum):
            image_tensor = image_and_label_obj['image_tensor']
            my_label = image_and_label_obj['label_tensor']
            w, h = image_tensor.shape[2:]
            # print(f'processing image size: {w}x{h}')

            # 原图展示
            image_basic = image_tensor[0].permute(1, 2, 0)
            # print(image_basic.shape)
            # plt.figure(figsize=(w / 100 + 1, h / 100 + 1))
            # plt.title('Original Image')
            # plt.imshow(image_basic.cpu().numpy())
            # plt.show()

            # 标签图展示
            label_basic = my_label
            # plt.figure(figsize=(w / 100 + 1, h / 100 + 1))
            # plt.title('Label Image')
            # plt.imshow(label_basic)
            # plt.show()
            #  直接使用predict_model测试数据   但是测试图片的格式必须是tensor类型 即 [batch,channel,height,width]
            pred_img_6_tensor = my_predict_model(image_tensor)
            #  取结果中6个尺度的第一个尺度图
            pred_img = pred_img_6_tensor[0]
            # print(pred_img[0].shape) #torch.Size([1, 256, 256])

            # 压缩维度(即将1维度去掉)并应用sigmoid
            pred_img_sigmoid = torch.sigmoid(pred_img.cpu().squeeze())

            # 转换为0-255范围的灰度图
            save_result_pred = (pred_img_sigmoid.data.numpy() * 255).astype(np.uint8)
            # print(pred_img_sigmoid.shape)
            # 显示预测图,尺寸比原图 大100px  用于边界信息显示
            # plt.figure(figsize=(w/100+1, h/100+1))
            # plt.title('Prediction Image')
            # plt.imshow(save_result_pred, cmap='gray')
            # plt.show()


            # 将Original Image、Label Image、Prediction Image 放在一张图上显示
            # 假设 original_img, label_img, pred_img 是 NumPy 数组或张量（形状为 [H, W, C]）
            fig, axes = plt.subplots(1, 3, figsize=(15, 6))  # 1行3列，横向排列

            # 显示 Original Image
            axes[0].imshow(image_basic.cpu().numpy())
            axes[0].set_title("Original Image")
            axes[0].axis('off')  # 关闭坐标轴

            # 显示 Label Image
            axes[1].imshow(label_basic)
            axes[1].set_title("Label Image")
            axes[1].axis('off')

            # 显示 Prediction Image
            axes[2].imshow(save_result_pred,cmap='gray')
            axes[2].set_title("Prediction Image")
            axes[2].axis('off')


            plt.tight_layout()  # 自动调整子图间距
            plt.show()



            # 计算指标
            label_basic = cv2.cvtColor(label_basic, cv2.COLOR_RGB2GRAY)  # 转换为灰度图  去掉颜色通道值 为了与预测图进行比较
            scaler = MinMaxScaler(feature_range=(0, 1))
            label_basic_normalized = scaler.fit_transform(label_basic.reshape(-1, 1))
            # print(label_basic.shape, save_result_pred.shape)
            batch_metrics = calculate_metrics(label_basic_normalized, save_result_pred)
            # print(batch_metrics)
            metrics.append(batch_metrics)
        print("Metrics:", metrics)